program Main

    use omp_lib
    use Globals
    use ModuleSolveEqm
    use ModuleQTRANSITION

    implicit none
    
    ! Local variable declarations
    real(8) :: param_vec_init(3)    ! vector of tax function parameters (initial system): input for function solving for equilibrium given taxes
    real(8) :: param_vec_lr(3)      ! vector of tax function parameters (final system): input for function solving for equilibrium given taxes
    real(8) :: welfare              ! auxiliary variable to store welfare numbers
    integer :: clck_counts_beg, clck_counts_end, clck_rate  ! for timing (true time; not CPU time with parallelization)
    real(8) :: ParamCalib(5)        ! calibrated parameters
    real(8) :: eqm_results(5)       ! equilibrium values from calibration
    real(8) :: V_TR_init(NS)        ! value function in first period of transition
    
    ! Start timer
    call system_clock ( clck_counts_beg, clck_rate )
    
    ! Choose normalization version
    norm_version = 1            ! indicator for choice of normalization: 1) mean income; 2) median income; 3) mean income of the third quintile
     
    ! Calibrated parameters
    open(11,  file = 'Input/ParamCalib.txt',      status = 'unknown')
    read(11, *) ParamCalib
    close(11)
    beta        = ParamCalib(1)             ! discount factor
    B           = ParamCalib(2)             ! labor disutility parameter
    G           = ParamCalib(3)             ! government spending
    Debt        = ParamCalib(4)             ! government debt
    ymean       = ParamCalib(5)             ! mean income
    ymeanCalib  = ParamCalib(5)             ! mean income
    
    ! Set tax function parameters
    open(11,  file = 'Input/eqm_results.txt',      status = 'unknown')
    read(11, *) eqm_results
    close(11)
    param_vec_init(1) = eqm_results(3)     ! gamma
    param_vec_init(2) = eqm_results(2)     ! lambda
    param_vec_init(3) = eqm_results(5)     ! rt
    omegat = 0d0                           ! omegat
    print *, 'Tax parameters initially: gamma, lambda, rt = ', param_vec_init
    print *, ''
    
    ! Guesses for equilibrium objects
    mt     = eqm_results(4)             ! mt
    r      = eqm_results(1)             ! interest rate
    lambda = param_vec_init(2)          ! lambda
    print*, 'Guesses for r = ', r,' and lambda = ', lambda
    print *, ''

    ! Solve for initial steady state
    print *, 'Computing initial steady-state'
    print *, ''
    
    ! Solve for equilibrium given tax function parameters
    welfare = welfare_eqm(param_vec_init,1)

    ! Store eqm objects
    r_init = r                  ! initial equilibrium interest rate
    mt_init = mt                ! initial equilibrium level transfer
    error_vec_init = error_vec  ! initial equilibrium vector of errors
    Kagg_init = Kagg            ! initial equilibrium capital stock

    ! Solve for new steady state
    print *, ''
    print *, 'Computing new steady-state'
    print *, ''

    ! Set tax function parameters for final equilibrium
    param_vec_lr(1) = 0d0       ! gamma
    param_vec_lr(2) = 0.605d0   ! lambda
    param_vec_lr(3) = 0d0       ! rt

    ! Solve for equilibrium given tax function parameters
    welfare = welfare_eqm(param_vec_lr,2)

    ! Store eqm objects
    r_lr = r                    ! new equilibrium interest rate
    lambda_lr = lambda          ! new equilibrium tax function level
    mt_lr = mt                  ! new equilibrium level transfer
    error_vec_lr  = error_vec   ! new equilibrium vector of errors

    ! Solving for transition
    print *, ''
    print *, 'Start transition'
    print *, ''
    
    ! Guesses for transition
    r_TR = r_lr                                     ! interest rate path
    lambda_TR = lambda_lr                           ! tax level path
    mt_TR = mt_lr                                   ! transfer level path
    wge_TR = (1.0D0-alpha)/(r_TR+delta)             ! wage path
    wge_TR = alpha*(wge_TR**((1d0-alpha)/alpha))    ! wage path

    ! Compute Jacobian
    call Compute_JACOB

    ! Compute transition
    call Compute_QNEWTON

    ! Welfare accounting for transition to new steady state
    welfare = -sum(mu_TR(:,1)*V_TR(:,1))     
    print *, 'Welfare transition = ', welfare
    print *, ''

    ! Save some output for producing figures
    ! Value function in first transition period
    V_TR_init = V_TR(:,1)
    open(1,  file = 'Output/V_TR.txt',  status = 'unknown')
    write(1, *) V_TR_init
    close(1)

    ! Transition matrix productivity
    open(2,  file = 'Output/Px.txt',      status = 'unknown')
    write(2, *) Px
    close(2)
    
    ! Asset policy along the transition
    open(20,  file = 'Output/a_pol_TR.txt',      status = 'unknown')
    write(20, *) a_pol_TR
    close(20)
    
    ! Consumption policy along the transition
    open(21,  file = 'Output/c_pol_TR.txt',      status = 'unknown')
    write(21, *) c_pol_TR
    close(21)

    ! Labor policy along the transition
    open(22,  file = 'Output/h_pol_TR.txt',      status = 'unknown')
    write(22, *) h_pol_TR
    close(22)

    ! End timer
    call system_clock ( clck_counts_end, clck_rate )
    print *, 'Time in seconds = ' , (clck_counts_end - clck_counts_beg) / real (clck_rate)

end program Main
